import os

from cryptography.fernet import Fernet


def generate_key(key_file_path):
    """
    生成密钥并保存到文件
    """
    key = Fernet.generate_key()
    with open(key_file_path, 'wb') as key_file:
        key_file.write(key)
    print(f"密钥已生成并保存到 '{key_file_path}'")


def encrypt_data(data, key_file_path, output_file_path):
    """
    加密数据并保存到文件
    """
    # 从密钥文件中读取密钥
    with open(key_file_path, 'rb') as key_file:
        key = key_file.read()

    cipher = Fernet(key)
    encrypted_data = cipher.encrypt(data.encode())

    # 保存加密后的数据
    with open(output_file_path, 'wb') as enc_file:
        enc_file.write(encrypted_data)

    print(f"数据已加密并保存到 '{output_file_path}'")


def decrypt_data(encrypted_file_name, key_file_name):
    """
    解密文件中的数据
    """
    # 从密钥文件中读取密钥、从自己文件目录下获取文件
    current_directory = os.path.dirname(os.path.abspath(__file__))
    encrypted_file_path = os.path.join(current_directory, encrypted_file_name)
    key_file_path = os.path.join(current_directory, key_file_name)

    with open(key_file_path, 'rb') as key_file:
        key = key_file.read()

    cipher = Fernet(key)

    # 从加密文件中读取加密数据
    with open(encrypted_file_path, 'rb') as enc_file:
        encrypted_data = enc_file.read()

    # 解密数据
    decrypted_data = cipher.decrypt(encrypted_data).decode()

    return decrypted_data


def main():
    # 生成密钥（仅需运行一次）
   # generate_key('secret.key')

    # 加密数据
    db_url = "mysql+pymysql://root:khZHuFbXz46YDwHvwP@rm-2vc0gfc166djy8f86.mysql.cn-chengdu.rds.aliyuncs.com:3306"
    encrypt_data(db_url, 'secret.key', 'encrypted_data_test.txt')

    # 解密数据
    decrypted_db_url = decrypt_data('encrypted_data_test.txt', 'secret.key')
    print("解密后的数据库 URL:", decrypted_db_url)

if __name__ == "__main__":
    main()